import logging
import json
from util.read_txt_file import ReadTxt
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
import os
from dotenv import load_dotenv

logger = logging.getLogger(__name__)

class BedSupplementaryCatalog:
    """年龄辅助目录处理类"""
    
    def __init__(self):
        self.bed_knowledge = self.load_bed_knowledge()
        
        # 初始化AI模型
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_BASE_URL")
        model_name = os.getenv("MODEL_NAME")
        
        self.llm = ChatOpenAI(
            temperature=0.1,
            max_tokens=300,
            base_url=base_url,
            model=model_name,
            api_key=api_key
        )
        
        # 创建床日分析提示词模板
        self.bed_prompt = ChatPromptTemplate.from_template(
            """你是DIP床日辅助分组专家，请根据患者床日信息和医疗知识进行精确的床日辅助分组调整。

患者信息：
- 患者姓名：{patient_name}
- 住院天数：{hospitalization_days}
- 诊断编码：{diagnosis_code}

原始AI分组结果：
{original_result}

床日辅助目录数据：
{bed_catalog_data}

床日医疗知识库：
{bed_knowledge}

请根据以上信息，判断是否需要根据住院天数进行分组调整，并给出调整后的结果。

调整规则：
1. 如果住院天数明显超出该疾病的正常范围，需要进行床日调整
2. 参考床日辅助目录中的标准住院天数
3. 考虑患者的具体诊断和治疗复杂程度
4. **重要：分值需要乘以住院天数进行计算，公式为：调整后分值 = 分值 × 住院天数**

请返回JSON格式的结果：

{{
  "分组名称": "调整后的治疗方式名称",
  "分值": 调整后的数值,
  "调整说明": "调整原因和计算过程",
  "床日调整完成": true
}}

**重要要求**：
1. 必须返回有效的JSON格式
2. 分值必须是数值类型，不是字符串
3. 不要添加任何JSON之外的文字

返回JSON："""
        )
        
        self.json_parser = JsonOutputParser()
        self.bed_chain = self.bed_prompt | self.llm | self.json_parser

    #加载年龄医疗知识库
    def load_bed_knowledge(self):

        try:
            reader_txt = ReadTxt('床日医疗知识库.txt')
            return reader_txt.load_txt_lines()
        except Exception as e:
            logger.error(f"加载床日医疗知识库失败: {str(e)}")
            return ""
    
    def analyze_bed_adjustment(self, patient_row, original_result,bed_catalog_data):

        try:
            result = self.bed_chain.invoke({
                "patient_name": patient_row['人员姓名'],
                "hospitalization_days": patient_row['住院天数'],
                "diagnosis_code" : patient_row['诊断编码'],
                "original_result": original_result,
                "bed_catalog_data": bed_catalog_data,
                "bed_knowledge": self.bed_knowledge
            })
            logger.info(f"床日智能体分析结果: {result}")
            
            # 如果返回的是字典，直接返回；如果是字符串，尝试解析JSON
            if isinstance(result, dict):
                return result
            elif isinstance(result, str):
                try:
                    return json.loads(result)
                except json.JSONDecodeError:
                    logger.error(f"床日智能体返回的不是有效JSON: {result}")
                    return {"错误": "JSON解析失败", "原始结果": result}
            else:
                return {"错误": "未知返回格式", "原始结果": str(result)}
                
        except Exception as e:
            logger.error(f"床日智能体分析失败: {str(e)}")
            return {"错误": f"床日调整失败: {str(e)}"}
    

